package gdma.hqdes.stat.renderer;

import org.geotools.data.DataUtilities;
import org.geotools.data.collection.ListFeatureCollection;
import org.geotools.data.geojson.GeoJSONDataStore;
import org.geotools.data.simple.SimpleFeatureCollection;
import org.geotools.data.simple.SimpleFeatureIterator;
import org.geotools.data.simple.SimpleFeatureSource;
import org.geotools.data.store.ContentDataStore;
import org.geotools.factory.CommonFactoryFinder;
import org.geotools.feature.simple.SimpleFeatureBuilder;
import org.geotools.feature.simple.SimpleFeatureTypeBuilder;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.map.FeatureLayer;
import org.geotools.map.Layer;
import org.geotools.map.MapContent;
import org.geotools.renderer.GTRenderer;
import org.geotools.renderer.lite.StreamingRenderer;
import org.geotools.styling.Stroke;
import org.geotools.styling.*;
import org.opengis.feature.simple.SimpleFeature;
import org.opengis.feature.simple.SimpleFeatureType;
import org.opengis.filter.FilterFactory2;
import org.springframework.stereotype.Service;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

/**
 * @Author wangyue
 * @Version 3.0
 * @Date 2024-08-28
 */
@Service
public class GeoMapRenderer {

    public static String NAME_KEY = "name";
    public static String VALUE_KEY = "value";

    public static Function<Number, Color> DEFAULT_COLOR_MAPPER = value -> {
        if (value.intValue() <= 50) {
            // 绿色到黄色的渐变
            float ratio = value.floatValue() / 50;
            int red = (int) (255 * ratio);
            int green = 255;
            return new Color(red, green, 0);
        } else {
            // 黄色到红色的渐变
            float ratio = (value.floatValue() - 50) / 50;
            int red = 255;
            int green = (int) (255 * (1 - ratio));
            return new Color(red, green, 0);
        }
    };

    public  <T extends Number> InputStream renderImage(String shortName, Function<T, Color> colorMapper,
                                                       Map<String, T> data, T zeroValue,
                                                       int width, int height) throws IOException {
        ContentDataStore dataStore = new GeoJSONDataStore(GeoMapRenderer.class.getResource("/geojson/" + shortName + ".json"));

        // 获取所有图层名称
        String typeName = dataStore.getTypeNames()[0]; // 假设我们只需要第一个图层

        // 获取指定图层的 FeatureSource
        SimpleFeatureSource featureSource = dataStore.getFeatureSource(typeName);

        // 创建新的 FeatureType，添加值属性
        SimpleFeatureType schema = featureSource.getSchema();
        SimpleFeatureTypeBuilder typeBuilder = new SimpleFeatureTypeBuilder();
        typeBuilder.init(schema);
        typeBuilder.add(VALUE_KEY, Integer.class);
        SimpleFeatureType newFeatureType = typeBuilder.buildFeatureType();

        // 创建新的 FeatureCollection，包含值数据
        SimpleFeatureCollection featureCollection = featureSource.getFeatures();
        SimpleFeatureBuilder featureBuilder = new SimpleFeatureBuilder(newFeatureType);
        List<SimpleFeature> newFeatures = new ArrayList<>();

        try (SimpleFeatureIterator iterator = featureCollection.features()) {
            while (iterator.hasNext()) {
                SimpleFeature feature = iterator.next();
                featureBuilder.init(feature);
                String name = (String) feature.getAttribute(NAME_KEY); // 假设省份名称属性为 "name"
                T value = data.getOrDefault(name, zeroValue);
                featureBuilder.set(VALUE_KEY, value);
                newFeatures.add(featureBuilder.buildFeature(feature.getID()));
            }
        }

        // 创建新的 FeatureSource
        SimpleFeatureCollection newFeatureCollection = new ListFeatureCollection(newFeatureType, newFeatures);
        SimpleFeatureSource newFeatureSource = DataUtilities.source(newFeatureCollection);

        // 创建地图内容
        MapContent map = new MapContent();

        // 创建 StyleBuilder 实例
        StyleBuilder styleBuilder = new StyleBuilder();

        // 定义 FilterFactory2 实例
        FilterFactory2 ff = CommonFactoryFinder.getFilterFactory2();

        // 使用 colorMapper 函数来设置颜色
        List<Rule> rules = new ArrayList<>();
        for (SimpleFeature feature : newFeatures) {
            T value = (T) feature.getAttribute(VALUE_KEY);
            Color color = colorMapper.apply(value);
            Fill fill = styleBuilder.createFill(color);
            Stroke stroke = styleBuilder.createStroke(Color.BLACK, 1);
            PolygonSymbolizer symbolizer = styleBuilder.createPolygonSymbolizer(stroke, fill);
            Rule rule = styleBuilder.createRule(symbolizer);
            rule.setFilter(ff.id(feature.getIdentifier()));
            rules.add(rule);
        }

        // 创建样式并添加规则
        Style style = styleBuilder.createStyle();
        style.featureTypeStyles().add(styleBuilder.createFeatureTypeStyle(typeName, rules.toArray(new Rule[0])));
        Layer layer = new FeatureLayer(newFeatureSource, style);

        map.addLayer(layer);

        // 创建渲染器
        GTRenderer renderer = new StreamingRenderer();
        renderer.setMapContent(map);

        // 创建图像
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
        Graphics2D graphics = image.createGraphics();
        graphics.setComposite(AlphaComposite.Clear);
        graphics.fillRect(0, 0, width, height);
        graphics.setComposite(AlphaComposite.SrcOver);

        // 获取地图边界
        ReferencedEnvelope mapBounds = newFeatureSource.getBounds();

        // 渲染地图到图像
        renderer.paint(graphics, new Rectangle(width, height), mapBounds);

        // 返回图像数据流
        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            ImageIO.write(image, "png", out);
            return new ByteArrayInputStream(out.toByteArray());
        } catch (Exception x) {
            throw new RuntimeException(x);
        }
    }

}
